# Written by Bram Cohen
# see LICENSE.txt for license information

from random import randrange, shuffle
from threading import RLock
from BitTornado.clock import clock

class Interests(object):
    """
    The interest levels.
    Each level is a list of piece indexes that are at that level
    """
    def __init__(self, priority_step, numpieces, lock):
        """ Save parameters and Initialize """
        # Parameters
        self.lock = lock
        self.numpieces = numpieces
        self.priority_step = priority_step
        # Initialize
        self.init()

    def init(self):
        """ Initialize the Interests object """
        
        self.lock.acquire()
        
        # Interests
        self.body = [[] for x in xrange(self.priority_step)]
        interests = range(self.numpieces)
        shuffle(interests)
        
        # Level in interests (self.body[X][])
        self.level_in_interests = [self.priority_step] * self.numpieces
        
        # Position in interests (self.body[][X])
        self.pos_in_interests = [0] * self.numpieces
        for i in xrange(self.numpieces):
            self.pos_in_interests[interests[i]] = i

        # Currently all pieces are in self.body[priority_step]
        self.body.append(interests)
        
        self.lock.release()
        
    def getLevel(self, piece):
        """ Get the level of the piece in the interests list """
        return self.level_in_interests[piece]

    def getPosition(self, piece):
        """ Get the position of the piece in its level """
        return self.pos_in_interests[piece]

    def add(self, piece, newLevel):
        """ Add a piece """
        self.lock.acquire()
        
        # Create this level if necessary
        while(len(self.body) <= newLevel):
            self.body.append([])

        # Find new position for our piece
        newPosition = randrange(len(self.body[newLevel])+1)
        if newPosition == len(self.body[newLevel]):
            # Add to the end
            self.body[newLevel].append(piece)
        else:
            # Take the piece in our position and put it in the end
            oldpiece = self.body[newLevel][newPosition]
            self.pos_in_interests[oldpiece] = len(self.body[newLevel])
            self.body[newLevel].append(oldpiece)
            # Put the piece in the right place
            self.body[newLevel][newPosition] = piece
        self.pos_in_interests[piece] = newPosition
        self.level_in_interests[piece] = newLevel
        
        self.lock.release()
        
    def remove(self, piece):
        """ Remove the piece """
        self.lock.acquire()
        
        # Retrieve level and position
        level = self.getLevel(piece)
        position = self.getPosition(piece)
        # Make sure it's the right piece
        assert self.body[level][position] == piece
        # Take the last piece in the level and put it in our piece's position
        last = self.body[level][-1]
        self.body[level][position] = last
        self.pos_in_interests[last] = position
        del self.body[level][-1]
        
        self.lock.release()
        
    def bump(self, piece):
        """ Move a piece to the end of its interest level """
        self.lock.acquire()
        
        level = self.getLevel(piece)
        position = self.getPosition(piece)
        assert self.body[level][position] == piece
        levelList = self.body[level]
        del levelList[position]
        levelList.append(piece)
        for i in xrange(position, len(levelList)):
            self.pos_in_interests[levelList[i]] = i
        
        self.lock.release()
                
    def increase(self, piece):
        """ change level to level + 1 """
        self.changeLevel(piece, self.getLevel(piece) + 1)
        
    def decrease(self, piece):
        """ change level to level - 1 """
        self.changeLevel(piece, self.getLevel(piece) - 1)

    def changeLevel(self, piece, newLevel):
        """ change the level of the piece """
        
        self.lock.acquire()
        
        # Retrieve old level
        oldLevel = self.getLevel(piece)

        # Retrieve old position
        oldPosition = self.getPosition(piece)

        # Make sure it's the right piece
        assert self.body[oldLevel][oldPosition] == piece

        # Create this level if necessary
        while(len(self.body) <= newLevel):
            self.body.append([])

        # Take the last piece in the old level and put it in our piece's position
        last = self.body[oldLevel][-1]
        self.body[oldLevel][oldPosition] = last
        self.pos_in_interests[last] = oldPosition
        del self.body[oldLevel][-1]

        # Find new position for our piece
        newPosition = randrange(len(self.body[newLevel])+1)
        if newPosition == len(self.body[newLevel]):
            # Add to the end
            self.body[newLevel].append(piece)
        else:
            # Take the piece in our position and put it in the end
            oldpiece = self.body[newLevel][newPosition]
            self.pos_in_interests[oldpiece] = len(self.body[newLevel])
            self.body[newLevel].append(oldpiece)
            # Put the piece in the right place
            self.body[newLevel][newPosition] = piece
        self.pos_in_interests[piece] = newPosition
        self.level_in_interests[piece] = newLevel
        
        self.lock.release()

    def decreaseAll(self):
        """ decrease the level of all pieces by one """
        self.lock.acquire()
        self.level_in_interests = [i-1 for i in self.level_in_interests]
        if self.body:
            del self.body[0]
        self.lock.release()

    def __len__(self):
        return len(self.body)
    
class PiecePicker:
    """
    Choose pieces to download
    """
    def __init__(self, numpieces, rarest_first_cutoff = 1, rarest_first_priority_cutoff = 3, priority_step = 20):        
        # number of completed piece downloads at which to switch from random piece selection to rarest first
        self.rarest_first_cutoff = rarest_first_cutoff
        # number of peers which need to have a piece before other partials take priority over rarest first
        self.rarest_first_priority_cutoff = rarest_first_priority_cutoff + priority_step
        # the difference bewteen priority levels 0, 1, and 2
        self.priority_step = priority_step
        # number of peers which need to have a piece before other partials take priority over rarest first
        self.cutoff = rarest_first_priority_cutoff
        # total number of pieces in the download
        self.numpieces = numpieces
        # the pieces that have been requested for download
        self.started = []
        # the total number of copies of all pieces in non-seeding peers
        self.totalcount = 0
        # the number of copies of each piece in non-seeding peers
        self.numhaves = [0] * numpieces
        # the priority of each piece (-1=don't download, 0=highest, 1=medium, 2=lowest)
        self.priority = [1] * numpieces
        # keys are the pieces that were started but then disabled
        self.removed_partials = {}
        # how many pieces X peers have
        self.crosscount = [numpieces]
        # like crosscount but includes this peer
        self.crosscount2 = [numpieces]
        # values are 1 for the pieces that this peer has
        self.has = [0] * numpieces
        # the number of pieces this peer has
        self.numgot = 0
        # whether the download is complete
        self.done = False
        # connections that have been made in super-seed mode
        self.seed_connections = {}
        # the IPs that have previously connected in super-seed mode
        self.past_ips = {}
        # the time when the first peer was seen in super-seed mode
        self.seed_time = None
        # whether we are in super-seed mode
        self.superseed = False
        # the number of connected seeds
        self.seeds_connected = 0
        # lock object
        self.lock = RLock()
        # Interests object
        self.interests = Interests(self.priority_step, self.numpieces, self.lock)

    def got_have(self, piece):
        """ Peer reported to have piece """
        # one more piece ;)
        self.totalcount += 1
        
        # update numhaves
        numint = self.numhaves[piece]
        self.numhaves[piece] += 1
        
        # update crosscount
        self.crosscount[numint] -= 1
        if numint+1==len(self.crosscount):
            self.crosscount.append(0)
        self.crosscount[numint+1] += 1

        # Still Downloading
        if not self.done:
            # update crosscount2
            numintplus = numint+self.has[piece]
            self.crosscount2[numintplus] -= 1
            if numintplus+1 == len(self.crosscount2):
                self.crosscount2.append(0)
            self.crosscount2[numintplus+1] += 1
            # if this piece is wanted - increase its level
            if not self.has[piece] and not self.is_blocked(piece):
                self.interests.increase(piece)
        # Super-seeding
        elif self.superseed:
            # update info for super-seed   
            self.seed_got_haves[piece] += 1
            # increase this piece's level
            self.interests.increase(piece)

    def lost_have(self, piece):
        """ Peer with a given piece disconnected """

        # one piece less ;(
        self.totalcount-=1

        # update numhaves
        numint = self.numhaves[piece]
        self.numhaves[piece] -= 1
        
        # update crosscount
        self.crosscount[numint] -= 1
        self.crosscount[numint-1] += 1

        # Still Downloading
        if not self.done:
            # update crosscount2
            numintplus = numint+self.has[piece]
            self.crosscount2[numintplus] -= 1
            self.crosscount2[numintplus-1] += 1
            # if this piece is wanted - decrease its level
            if not self.has[piece] and not self.is_blocked(piece):
                self.interests.decrease(piece)
        # Super-seeding
        elif self.superseed:
            self.interests.decrease(piece)

    def got_seed(self):
        """ Found a new seed """
        self.seeds_connected += 1
        self.cutoff = max(self.rarest_first_priority_cutoff-self.seeds_connected,0)

    def became_seed(self):
        """ A known peer became a seed """
        self.got_seed()
        self.totalcount -= self.numpieces
        self.numhaves = [i-1 for i in self.numhaves]
        if self.superseed or not self.done:
            self.interests.decreaseAll()
        # TODO: understand this one :|
        del self.crosscount[0]
        if not self.done:
            del self.crosscount2[0]

    def lost_seed(self):
        """ A seed disconnected """
        self.seeds_connected -= 1
        self.cutoff = max(self.rarest_first_priority_cutoff-self.seeds_connected,0)

    def requested(self, piece):
        """ Add a piece that a request has been sent for to the list of started pieces """
        if piece not in self.started:
            self.started.append(piece)

    def _remove_from_interests(self, piece, keep_partial = False):
        """ Remove a piece from the interests """
        self.interests.remove(piece)
        try:
            self.started.remove(piece)
            if keep_partial:
                self.removed_partials[piece] = 1
        except ValueError:
            pass

    def complete(self, piece):
        """ Successfully received piece """
        assert not self.has[piece]
        self.has[piece] = 1
        self.numgot += 1
        if self.numgot == self.numpieces:
            self.done = True
            self.crosscount2 = self.crosscount
        else:
            numhaves = self.numhaves[piece]
            self.crosscount2[numhaves] -= 1
            if numhaves+1 == len(self.crosscount2):
                self.crosscount2.append(0)
            self.crosscount2[numhaves+1] += 1
        self._remove_from_interests(piece)


    def next(self, haves, wantfunc, complete_first = False):
        """ Choose a piece to request to download next from a peer """

        cutoff = self.numgot < self.rarest_first_cutoff
        complete_first = (complete_first or cutoff) and not haves.complete()
        best = None
        bestnum = 2 ** 30
        for i in self.started:
            if haves[i] and wantfunc(i):
                level = self.interests.getLevel(i)
                if level < bestnum:
                    best = i
                    bestnum = level
        if best is not None:
            if complete_first or (cutoff and len(self.interests) > self.cutoff):
                return best
        if haves.complete():
            r = [ (0, min(bestnum, len(self.interests))) ]
        elif cutoff and len(self.interests) > self.cutoff:
            r = [ (self.cutoff, min(bestnum, len(self.interests))), (0, self.cutoff) ]
        else:
            r = [ (0, min(bestnum, len(self.interests))) ]
            
        for lo, hi in r:
            for i in xrange(lo, hi):
                for j in self.interests.body[i]:
                    if haves[j] and wantfunc(j):
                        return j
        if best is not None:
            return best
        return None

    def am_I_complete(self):
        """ Determine if the download is complete """
        return self.done
    
    def bump(self, piece):
        """ Move a piece to the end of its interest level """
        self.interests.bump(piece)
        try:
            self.started.remove(piece)
        except:
            pass

    def set_priority(self, piece, p):
        """ Adjust the priority for a piece """
        # Don't muck with this if you're a superseed
        if self.superseed:
            return False
        
        # Set new priority
        oldp = self.priority[piece]
        if oldp == p:
            return False
        self.priority[piece] = p
        
        # Changed priority to never download
        if p == -1:
            # Make sure to cancel any downloads for this piece
            if not self.has[piece]:
                self._remove_from_interests(piece, True)
            return True
        # Changed priority from never download to download
        if oldp == -1:
            # TODO: the original code changed the piece level BEFORE the following condition. why? 
            if self.has[piece]:
                return True
            level = self.numhaves[piece] + (self.priority_step * p)
            self.interests.add(piece, level)
            if self.removed_partials.has_key(piece):
                del self.removed_partials[piece]
                self.started.append(piece)
            # now go to downloader and try requesting more
            return True
        
        # Simple priority change
        # TODO: the original code changed the piece level BEFORE the following condition. why? 
        if self.has[piece]:
            return False
        newLevel = self.interests.getLevel(piece) + ((p - oldp) * self.priority_step)
        self.interests.changeLevel(piece, newLevel)

        return False

    def is_blocked(self, piece):
        """ Determine whether a piece is disabled """
        return self.priority[piece] < 0


    def set_superseed(self):
        """ Switch to super-seeding mode """
        assert self.done
        self.superseed = True
        self.seed_got_haves = [0] * self.numpieces
        self.interests.init()  # assume everyone is disconnected

    def next_have(self, connection, looser_upload):
        """ Determine the next piece to tell a peer we have in super-seed mode """
        if self.seed_time is None:
            self.seed_time = clock()
            return None
        if clock() < self.seed_time+10:  # wait 10 seconds after seeing the first peers
            return None                 # to give time to grab have lists
        if not connection.upload.super_seeding:
            return None
        olddl = self.seed_connections.get(connection)
        if olddl is None:
            ip = connection.get_ip()
            olddl = self.past_ips.get(ip)
            if olddl is not None:                           # peer reconnected
                self.seed_connections[connection] = olddl
                if not looser_upload:
                    self.seed_got_haves[olddl] -= 1         # penalize
        if olddl is not None:
            if looser_upload:
                num = 1     # send a new have even if it hasn't spread that piece elsewhere
            else:
                num = 2
            if self.seed_got_haves[olddl] < num:
                return None
            if not connection.upload.was_ever_interested:   # it never downloaded it?
                connection.upload.skipped_count += 1
                if connection.upload.skipped_count >= 3:    # probably another stealthed seed
                    return -1                               # signal to close it
        for tier in self.interests.body:
            for piece in tier:
                if not connection.download.have[piece]:
                    self.interests.increase(piece)          # tweak it up one, so you don't duplicate effort
                    self.seed_got_haves[piece] = 0          # reset this
                    self.seed_connections[connection] = piece
                    connection.upload.seed_have_list.append(piece)
                    return piece
        return -1       # something screwy; terminate connection

    def lost_peer(self, connection):
        """ Process a lost peer in super-seed mode """
        olddl = self.seed_connections.get(connection)
        if olddl is None:
            return
        del self.seed_connections[connection]
        self.past_ips[connection.get_ip()] = olddl
        if self.seed_got_haves[olddl] == 1:
            self.seed_got_haves[olddl] = 0

class FakePiecePicker:
    """
    Fake Piece Picker
    """
    def __init__(self, numpieces = 1, rarest_first_cutoff = 1, rarest_first_priority_cutoff = 3, priority_step = 20):        
        self.rarest_first_cutoff = rarest_first_cutoff
        self.rarest_first_priority_cutoff = rarest_first_priority_cutoff + priority_step
        self.priority_step = priority_step
        self.cutoff = rarest_first_priority_cutoff
        self.numpieces = numpieces
        self.started = []
        self.totalcount = 0
        self.numhaves = [0] * numpieces
        self.priority = [1] * numpieces
        self.removed_partials = {}
        self.crosscount = [numpieces]
        self.crosscount2 = [numpieces]
        self.has = [0] * numpieces
        self.numgot = 0
        self.done = False
        self.seed_connections = {}
        self.past_ips = {}
        self.seed_time = None
        self.superseed = False
        self.seeds_connected = 0
    def got_have(self, piece):
        pass    
    def lost_have(self, piece):
        pass
    def got_seed(self):
        self.seeds_connected += 1
    def became_seed(self):
        self.got_seed()
    def lost_seed(self):
        self.seeds_connected -= 1
    def requested(self, piece):
        pass
    def _remove_from_interests(self, piece, keep_partial = False):
        pass    
    def complete(self, piece):
        pass
    def next(self, haves, wantfunc, complete_first = False):
        return None
    def am_I_complete(self):
        return self.done    
    def bump(self, piece):
        pass
    def set_priority(self, piece, p):
        return False        
    def is_blocked(self, piece):
        return True
    def set_superseed(self):
        pass
    def next_have(self, connection, looser_upload):
        return None
    def lost_peer(self, connection):
        pass
    
